#include <cmath>

#include "mpc_quadrotor_mujoco/QuadrotorSystemDynamics.h"

namespace ocs2 {
    namespace quadrotor {
        vector_t QuadrotorSystemDynamics::computeFlowMap(scalar_t time, const vector_t& state, const vector_t& input, const PreComputation&) {

            Eigen::Matrix<scalar_t, 3, 1> eulerAngle = state.segment<3>(3);
            Eigen::Matrix<scalar_t, 3, 3> JRPY = getMppingFromLocalAngularVelocityToRPYDerivative(eulerAngle);

            scalar_t posX = state(0);
            scalar_t posY = state(1);
            scalar_t posZ = state(2);
            
            scalar_t phi = state(3);
            scalar_t theta = state(4);
            scalar_t psi = state(5);
            
            scalar_t dPosX = state(6);
            scalar_t dPosY = state(7);
            scalar_t dPosZ = state(8);
            
            scalar_t OmegaX = state(9);
            scalar_t OmegaY = state(10);
            scalar_t OmegaZ = state(11);

            Eigen::Matrix<scalar_t, 3, 1> dEuler = JRPY * state.segment<3>(9);

            scalar_t Fz = input(0);
            scalar_t Mx = input(1);
            scalar_t My = input(2);
            scalar_t Mz = input(3);
            
            vector_t stateDerivative(STATE_DIM);
            stateDerivative[0] = dPosX;
            stateDerivative[1] = dPosY;
            stateDerivative[2] = dPosZ;
            stateDerivative[3] = dEuler(0);
            stateDerivative[4] = dEuler(1);
            stateDerivative[5] = dEuler(2);

            stateDerivative[6] = Fz * (cos(phi) * sin(theta) * cos(psi) + sin(phi) * sin(psi)) / param_.quadrotorMass_;
            stateDerivative[7] = Fz * (cos(phi) * sin(theta) * sin(psi) - sin(phi) * cos(psi)) / param_.quadrotorMass_;
            stateDerivative[8] = Fz * cos(phi) * cos(theta) / param_.quadrotorMass_ - param_.gravity_;
            stateDerivative[9] = (Mx + OmegaY * OmegaZ * (param_.Thyy_ - param_.Thzz_)) / param_.Thxx_;
            stateDerivative[10] = (My + OmegaX * OmegaZ * (param_.Thzz_ - param_.Thxx_)) / param_.Thyy_;
            stateDerivative[11] = (Mz + OmegaX * OmegaY * (param_.Thxx_ - param_.Thyy_)) / param_.Thzz_;

            return stateDerivative;
        }
        
        VectorFunctionLinearApproximation QuadrotorSystemDynamics::linearApproximation(scalar_t t, const vector_t& x, const vector_t& u,
                                                                                        const PreComputation& preComp) {
            VectorFunctionLinearApproximation dynamics;
            dynamics.f = computeFlowMap(t, x, u, preComp);

            Eigen::Matrix<scalar_t, 3, 1> eulerAngle = x.segment<3>(3);
            Eigen::Matrix<scalar_t, 3, 1> angularVelocity = x.segment<3>(9);
            Eigen::Matrix<scalar_t, 3, 6> jacobianOfAngularVelocityMapping_ = JacobianOfAngularVelocityMapping(eulerAngle, angularVelocity);

            scalar_t posX = x(0);
            scalar_t posY = x(1);
            scalar_t posZ = x(2);

            scalar_t phi = x(3);
            scalar_t theta = x(4);
            scalar_t psi = x(5);

            scalar_t sx = sin(phi);
            scalar_t cx = cos(phi);
            scalar_t sy = sin(theta);
            scalar_t cy = cos(theta);
            scalar_t sz = sin(psi);
            scalar_t cz = cos(psi);

            scalar_t dPosX = x(6);
            scalar_t dPosY = x(7);
            scalar_t dPosZ = x(8);

            scalar_t omegaX = x(9);
            scalar_t omegaY = x(10);
            scalar_t omegaZ = x(10);

            scalar_t Fz = u(0);
            scalar_t Mx = u(0);
            scalar_t My = u(0);
            scalar_t Mz = u(0);
            
            {
                matrix_t& A = dynamics.dfdx;
                A.setZero(STATE_DIM, STATE_DIM);
                A.block<3, 3>(0, 6).setIdentity();
                A.block<3, 3>(3, 3) = jacobianOfAngularVelocityMapping_.block<3, 3>(0, 0);
                A.block<3, 3>(3, 9) = jacobianOfAngularVelocityMapping_.block<3, 3>(0, 3);

                A(6, 3) = (-sx * sy * cz + cx * sz) * Fz / param_.quadrotorMass_;
                A(6, 4) = (cx * cy * cz) * Fz / param_.quadrotorMass_;
                A(6, 5) = (-cx * sy * sz + sx * cz) * Fz / param_.quadrotorMass_;
                A(7, 3) = (-sx * sy * sz - cx * cz) * Fz / param_.quadrotorMass_;
                A(7, 4) = (cx * cy * sz) * Fz / param_.quadrotorMass_;
                A(7, 5) = (cx * sy * cz + sx * sz) * Fz / param_.quadrotorMass_;
                A(8, 3) = (-sx * cy) * Fz / param_.quadrotorMass_;
                A(8, 4) = (-cx * sy) * Fz / param_.quadrotorMass_;

                A(9, 10) = (omegaZ * (param_.Thyy_ - param_.Thzz_)) / param_.Thxx_;
                A(9, 11) = (omegaY * (param_.Thyy_ - param_.Thzz_)) / param_.Thxx_;
                A(10, 9) = (omegaZ * (param_.Thzz_ - param_.Thxx_)) / param_.Thyy_;
                A(10, 11) = (omegaX * (param_.Thzz_ - param_.Thxx_)) / param_.Thyy_;
                A(11, 9) = (omegaY * (param_.Thxx_ - param_.Thyy_)) / param_.Thzz_;
                A(11, 10) = (omegaX * (param_.Thxx_ - param_.Thyy_)) / param_.Thzz_;
            }

            {
                matrix_t& B = dynamics.dfdu;
                B.setZero(STATE_DIM, INPUT_DIM);
                B(6, 0) = (cx * sy * cz + sx * sz) / param_.quadrotorMass_;
                B(7, 0) = (cx * sy * sz - sx * cz) / param_.quadrotorMass_;
                B(8, 0) = (cx * cy) / param_.quadrotorMass_;
                B(9, 1) = 1.0 / param_.Thxx_;
                B(10, 2) = 1.0 / param_.Thyy_;
                B(11, 3) = 1.0 / param_.Thzz_;
            }
            return dynamics;
        }

    } // namespace quadrotor
} // namespace ocs2

